package com.xxx.springboot;

import java.io.File;
import java.io.IOException;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;

import org.apache.commons.io.FileUtils;

import redis.clients.jedis.Jedis;

/**
 * redis + lua 实现限流
 */
public class RedisLimitRateWithLUA {
	
	// 限流阈值
	static String LIMIT = "5";

	public static void main(String[] args) {
		
		// 定义一个CountDownLatch（闭锁，等待其他线程各自执行完毕后再执行）
		final CountDownLatch latch = new CountDownLatch(1);

		// 创建7个线程，并启动，模拟同时有多个请求
		for (int i=0; i<7; i++) {
			new Thread(new Runnable() {
				public void run() {
					try {
						latch.await(); // 当前线程开始执行前，等待其他线程（主线程）执行完毕
						System.out.println("请求是否被执行：" + accquire());
					} catch (Exception e) {
						e.printStackTrace();
					}
				}
			}).start();
		}

		latch.countDown(); // 计数器减1，直到为0，唤醒await处等待的线程（即for循环里创建的7个子线程）
	}

	private static boolean accquire() throws IOException, URISyntaxException {
		// jedis
		Jedis jedis = new Jedis("127.0.0.1");
		jedis.auth("123456");
		
		// 加载lua脚本文件
		File luaFile = new File(RedisLimitRateWithLUA.class.getResource("/").toURI().getPath() + "limit.lua");
		String luaScript = FileUtils.readFileToString(luaFile, "UTF-8");

		// lua脚本参数：keys、 args
		String key = "ip:" + System.currentTimeMillis() / 1000; // 当前秒
		List<String> keys = new ArrayList<String>();
		keys.add(key);
		List<String> args = new ArrayList<String>();
		args.add(LIMIT); // 限流次数：5次/秒
		
		// 执行lua脚本，传入参数
		Long result = (Long) (jedis.eval(luaScript, keys, args));
		return result == 1; // true：可以执行，false：被限流
	}
	
}
